#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
    Module Documentation
    here
"""

# Created by  : Zhang Chengdong
# Create Date : 2024/11/22 13:44
# Version = v0.1.0

__author__ = "Zhang Chengdong"
__copyright__ = "Copyright 2024. Large scale model"
__credits__ = ['Zhang Chengdong']

__liscence__ = "MIT"
__version__ = "1.0.1"
__maintainer__ = "Zhang Chengdong"
__status__ = "Production"

import json
import logging
import os

import joblib
import numpy as np
import pandas as pd
from datetime import datetime

import yaml
from pymongo import MongoClient
from service.data_preprocessing.data_processing import DataProcessing
from typing import List, Union, Dict
from collections import defaultdict
from pathlib import Path

from service.catboost_model.platform_api_handler import PlatformAPIHandler
from service.catboost_model.model_architecture import NonLinearModel as non_model
from service.data_preprocessing.data_processing import DataProcessing as data_processing


def get_master_data_reverse(master_data: dict) -> dict:
    """
    获取master_data的反转
    :return:
    """
    master_data_reverse = {value: key for key, value in master_data.items()}
    return master_data_reverse


def format_single_info_data(product_code: str, model_version: str, mas_model_type: str, mas_model_code: str,
                            model3_n1d: tuple, model3_1d: tuple, model28_n1d: tuple, model28_1d: tuple,
                            model28_3d: tuple, end_time: str, start_time: str):
    """
    插入数据库的中数据格式组合
    :param mas_model_code:
    :param product_code:
    :param model_version:
    :param mas_model_type
    :param model3_n1d:
    :param model3_1d:
    :param model28_n1d:
    :param model28_1d:
    :param model28_3d:
    :param end_time:
    :param start_time:
    """
    single_info_data = {
        "productCode": product_code,
        "createdByName": "api",
        "modelType": mas_model_type,
        "name": "水泥强度预测通用模型",
        "version": model_version,
        "modelCode": mas_model_code,
        "evaluation_3d_n1d": {
            "mse": model3_n1d[0],
            "r2": model3_n1d[1],
            "acc": model3_n1d[2]
        },
        "evaluation_3d_1d": {
            "mse": model3_1d[0],
            "r2": model3_1d[1],
            "acc": model3_1d[2]
        },
        "evaluation_28d_n1d": {
            "mse": model28_n1d[0],
            "r2": model28_n1d[1],
            "acc": model28_n1d[2]
        },
        "evaluation_28d_1d": {
            "mse": model28_1d[0],
            "r2": model28_1d[1],
            "acc": model28_1d[2]
        },
        "evaluation_28d_3d": {
            "mse": model28_3d[0],
            "r2": model28_3d[1],
            "acc": model28_3d[2]
        },
        "endDate": end_time,
        "startDate": start_time
    }
    return single_info_data


def insert_json_data(db_url: str, db_name: str, collection_name: str, json_data: dict):
    """
    向数据库中插入结果数据
    :param db_url:
    :param db_name:
    :param collection_name:
    :param json_data:
    """
    with MongoClient(db_url) as client:
        db = client[db_name]

        if collection_name not in db.list_collection_names():
            db.create_collection(collection_name)
            logging.info("{}表不存在，已创建".format(collection_name))

        if isinstance(json_data, dict):
            result = db[collection_name].insert_one(json_data)
            logging.info("将模型训练的数据添加到表中".format(result.inserted_id))


def delete_all_files_in_directory(directory:str):
    """
    删除路径下的所有文件信息
    """
    path = Path(directory)
    for file_path in path.glob("*"):  # 遍历目录中的所有文件
        try:
            if file_path.is_file() or file_path.is_symlink():  # 如果是文件或符号链接
                file_path.unlink()  # 删除文件
                print(f"Deleted: {file_path}")
            elif file_path.is_dir():  # 如果是子目录
                print(f"Skipping directory: {file_path}")
        except Exception as e:
            print(f"Failed to delete {file_path}. Reason: {e}")


def train_main(data: pd.DataFrame, save_model_path: str, model_name: str = "CatBoost", model_desc_path: str = None,
               mas_model_type: str = None, params_for_3d: dict = None, params_for_28d: dict = None,
               train_optuna: dict = None, model_type_suffix: dict = None,
               mongo_url: str = None, database_name: str = None, tenant_id_list: list = None,
               save_zip_file: str = None, query_api: str = None, upload_api: str = None, register_api: str = None,
               deploy_api: str = None):
    """
    训练模型的主方法
    :param data：所有数据
    :param save_model_path：模型文件保存路径
    :param model_name：catboost
    :param model_desc_path: 描述文件json
    :param mas_model_type：通用模型标识
    :param params_for_3d:
    :param params_for_28d:
    :param train_optuna:
    :param model_type_suffix:
    :param mongo_url:
    :param database_name:
    :param tenant_id_list:
    :param save_zip_file:
    :param query_api:
    :param upload_api:
    :param register_api:
    :param deploy_api:
    :return:
    """
    data_df = data.copy()
    data_groups = data_df.groupby("品种")
    get_map_dict = data_processing(mongo_url, database_name, tenant_id_list)
    for product_code, data_group in data_groups:
        # TODO 上线需要注释
        # if product_code != "M32.5":
        #     continue

        data_group.reset_index(inplace=True)
        data_group.rename(columns={"index":"index_0"}, inplace=True)
        tenant_id = data_group.loc[:, "tenant_id"].unique().tolist()[0]
        master_data = get_map_dict.get_dictionary_map("DICTIONARY_{}".format(tenant_id))
        master_data_reverse = get_master_data_reverse(master_data)

        start_time_dt = datetime.now()
        start_time_str = start_time_dt.strftime('%Y-%m-%d %H:%M')
        logging.info("开始训练模型，品种：{}，数据量：{}".format(product_code, len(data_group)))
        product_code_type = master_data_reverse[product_code]
        mas_model_code = mas_model_type + "_" + product_code_type
        model = non_model(data_group, save_model_path, model_name, mas_model_code, train_optuna, model_type_suffix)

        # 3天模型相关开关信息
        need_combination_3 = params_for_3d.get("need_combination")
        need_chemical_3 = params_for_3d.get("need_chemical")
        need_pca_3 = params_for_3d.get("need_pca")
        need_rfe_3 = params_for_3d.get("need_rfe")
        promotion_rate_3 = {"use_1": True, "use_3": False}

        # 28天模型相关开关信息
        need_combination_28 = params_for_28d.get("need_combination")
        need_chemical_28 = params_for_28d.get("need_chemical")
        need_pca_28 = params_for_28d.get("need_pca")
        need_rfe_28 = params_for_28d.get("need_rfe")
        promotion_rate_28 = params_for_28d.get("promotion_rate")

        # 训练3天模型，无一天
        model3_n1d_mse, model3_n1d_r2, model3_n1d_acc = model.train_model_3d(actual1d=False,
                                                                             need_combination=need_combination_3,
                                                                             need_chemical=need_chemical_3,
                                                                             need_pca=need_pca_3,
                                                                             need_rfe=need_rfe_3,
                                                                             promotion_rate=promotion_rate_3)
        # model3_n1d_mse, model3_n1d_r2 = None, None
        logging.info("训练3天模型，无一天，训练的模型信息：MSE：{}，R2：{}".format(str(model3_n1d_mse), str(model3_n1d_r2)))
        # 训练3天模型,有一天
        model3_1d_mse, model3_1d_r2, model3_1d_acc = model.train_model_3d(actual1d=True,
                                                                          need_combination=need_combination_3,
                                                                          need_chemical=need_chemical_3,
                                                                          need_pca=need_pca_3,
                                                                          need_rfe=need_rfe_3,
                                                                          promotion_rate=promotion_rate_3)
        # model3_1d_mse, model3_1d_r2 = None, None
        logging.info("训练3天模型,有一天，训练的模型信息：MSE：{}，R2：{}".format(str(model3_1d_mse), str(model3_1d_r2)))
        #
        # # 训练28天模型，无一天
        model28_n1d_mse, model28_n1d_r2, model28_n1d_acc = model.train_model_28d(actual1d=False, actual3d=False,
                                                                                 need_combination=need_combination_28,
                                                                                 need_chemical=need_chemical_28,
                                                                                 need_pca=need_pca_28,
                                                                                 need_rfe=need_rfe_28,
                                                                                 promotion_rate=promotion_rate_28)
        # model28_n1d_mse, model28_n1d_r2 = None, None
        logging.info(
            "训练28天模型，无一天，训练的模型信息：MSE：{}，R2：{}".format(str(model28_n1d_mse), str(model28_n1d_r2)))
        #
        # # 训练28天模型，有一天，无3天
        model28_1d_mse, model28_1d_r2, model28_1d_acc = model.train_model_28d(actual1d=True, actual3d=False,
                                                                              need_combination=need_combination_28,
                                                                              need_chemical=need_chemical_28,
                                                                              need_pca=need_pca_28,
                                                                              need_rfe=need_rfe_28,
                                                                              promotion_rate=promotion_rate_28)
        # model28_1d_mse, model28_1d_r2 = None, None
        logging.info(
            "训练28天模型，有一天，无3天，训练的模型信息：MSE：{}，R2：{}".format(str(model28_1d_mse), str(model28_1d_r2)))
        #
        # # 训练28天模型，有一天，有3天
        model28_3d_mse, model28_3d_r2, model28_3d_acc = model.train_model_28d(actual1d=False, actual3d=True,
                                                                              need_combination=need_combination_28,
                                                                              need_chemical=need_chemical_28,
                                                                              need_pca=need_pca_28,
                                                                              need_rfe=need_rfe_28,
                                                                              promotion_rate=promotion_rate_28)
        # model28_3d_mse, model28_3d_r2 = None, None
        logging.info(
            "训练28天模型，有一天，有3天，训练的模型信息：MSE：{}，R2：{}".format(str(model28_3d_mse), str(model28_3d_r2)))

        end_time_dt = datetime.now()
        end_time_str = end_time_dt.strftime('%Y-%m-%d %H:%M')
        # model_version = "test_version_1"
        # 打包上传模型到模型平台
        model_version = packaging_and_upload(model_code=mas_model_code, model_desc=model_desc_path,
                                             save_zip_path=save_model_path, save_zip_file=save_zip_file,
                                             query_api=query_api, upload_api=upload_api, register_api=register_api,
                                             deploy_api=deploy_api)  # TODO 如果上线需要将该部分进行打开


        model3_n1d, model3_1d, model28_n1d, model28_1d, model28_3d = (model3_n1d_mse, model3_n1d_r2, model3_n1d_acc), \
            (model3_1d_mse, model3_1d_r2, model3_1d_acc), \
            (model28_n1d_mse, model28_n1d_r2, model28_n1d_acc), \
            (model28_1d_mse, model28_1d_r2, model28_1d_acc), \
            (model28_3d_mse, model28_3d_r2, model28_3d_acc)

        delete_all_files_in_directory(os.path.join(save_model_path,"CatBoost"))

    return {'status': 'success', "message": "模型训练成功"}


def rewrite_model_desc(model_path, model_code):
    """
    重写model desc文件
    :param model_path:
    :param model_code:
    :return:
    """

    # 读取 JSON 文件
    with open(model_path, 'r', encoding="utf-8") as file:
        data = json.load(file)

    data['modelCode'] = model_code
    # 保存修改后的 JSON 数据到文件
    with open(model_path, 'w', encoding="utf-8") as file:
        json.dump(data, file, indent=4, ensure_ascii=False)


def packaging_and_upload(model_code: str, model_desc: str, save_zip_path: str, save_zip_file: str, query_api: str,
                         upload_api: str, register_api: str, deploy_api: str):
    """
    打包相关代码和上传到大模型
    :param model_code:
    :param model_desc:
    :param save_zip_path:
    :param save_zip_file:
    :param query_api:
    :param upload_api:
    :param register_api:
    :param deploy_api:
    :return:
    """
    new_path = os.path.join(save_zip_path, "CatBoost")
    zip_path = os.path.join(new_path, "model.zip")
    with open(os.path.join(new_path, "predict.py"), "w", encoding="utf-8") as file:
        file.write("# 通用模型(模型平台需要的文件)\n")

    PlatformAPIHandler.pack_to_zip(directory=new_path, zip_file=zip_path)
    logging.info("----zip打包完成")
    # 3. upload

    logging.info(f"----模型开始上传")
    flag, model_mini_version = PlatformAPIHandler.register_new_version_to_platform(
        directory_need_zip=save_zip_path,
        zip_file_path=save_zip_file,
        model_code=model_code,
        query_api=query_api,
        upload_api=upload_api,
        register_api=register_api,
        model_name=model_code,
    )
    # 4. local
    PlatformAPIHandler.develop_in_platform(model_code=model_code,
                                           develop_type="LOCAL",
                                           model_version=model_mini_version,
                                           api_url=deploy_api)

    # TODO 后期确认是否部署到生产环境
    # 5. dev
    # PlatformAPIHandler.develop_in_platform(model_code=model_code,
    #                                        develop_type="PROD",
    #                                        model_version=model_mini_version,
    #                                        api_url=settings.deploy_api)
    return model_mini_version


if __name__ == '__main__':
    data_all = pd.read_excel("../base_data/demo.xlsx")
    with open("../../config/config.yaml", "r", encoding="utf-8") as f:
        choice_config = yaml.safe_load(f)

    train_main(data=data_all, save_model_path="./model_path", model_name=choice_config.get("model_name"),
               mas_model_type=choice_config.get("mas_model_type"), params_for_3d=choice_config.get("params_for_3d"),
               params_for_28d=choice_config.get("params_for_28d"),
               train_optuna=choice_config.get("train_optuna"),
               model_type_suffix=choice_config.get("model_type_suffix"),
               mongo_url=choice_config.get("testing_environment").get("mongo_uri"),
               database_name=choice_config.get("testing_environment").get("database_name"),
               tenant_id_list=[100])
